import sys
import torch
import torch_npu

npu_id=sys.argv[1]
# 在NPU设备上创建测试数据
device = torch.device(f"npu:{npu_id}")
input = torch.randn(100, 100, device=device)
mask = torch.rand(100, 100, device=device) > 0.5

# 执行MaskedSelect
try:
    output = torch.masked_select(input, mask)
    print(output)
except Exception as e:
    print(f"Error occurred: {e}")
